from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, File, UploadFile
import cv2
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse, StreamingResponse
from ultralytics import YOLO


cap = None  # 全局摄像头实例
model = YOLO("best.pt")     # 动物类别检测模型


@asynccontextmanager
async def lifespan(app: FastAPI):
    """管理应用生命周期"""

    global cap

    # 初始化摄像头
    camera_number = 9
    gst_pipeline = (
        f"v4l2src device=/dev/video{camera_number} ! "
        "image/jpeg, width=640, height=480, framerate=15/1 ! "
        "jpegdec ! videoconvert ! appsink"
    )
    cap = cv2.VideoCapture(gst_pipeline, cv2.CAP_GSTREAMER)
    # cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise RuntimeError("无法打开摄像头")


    # 初始化压力传感器
    # todo: 未实现的功能

    # 启动传感器监听线程

    # 启动生命周期
    yield

    # 关闭资源(摄像头和传感器)
    if cap:
        cap.release()
    # GPIO.cleanup()
    print("资源已释放")


app = FastAPI(lifespan=lifespan)

# 允许跨域访问
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 设置为 "*" 允许所有来源访问
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=True,
    expose_headers = {"Content-Disposition"},

)


@app.get("/detect_video", summary='动物检测检测视频流')
def detect_video():
    # 返回目标检测视频流
    global cap
    if not cap or not cap.isOpened():
        return JSONResponse({'message': '请先打开摄像头'}, status_code=400)

    def get_video_stream():
        global cap, model
        if not cap.isOpened():
            raise RuntimeError("无法打开摄像头")

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # YOLO 检测
            results = model(frame)  # 进行检测

            # 绘制结果
            frame = results[0].plot()
            # 将帧编码为 JPEG 格式
            _, jpeg = cv2.imencode(".jpg", frame)
            encoded_frame = jpeg.tobytes()

            # 发送帧数据
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + encoded_frame + b'\r\n')

        # cap.release()

    return StreamingResponse(get_video_stream(), media_type="multipart/x-mixed-replace; boundary=frame")


@app.get("/")
async def root():
    return {"message": "压力传感器监控运行中"}


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host='0.0.0.0', port=18888)
